# -*- coding: utf-8 -*-
"""
File: url_pool.py
Author: StellaJL
Date: 2025-09-16 01-42-04
Version：1.0.0
Description: 引擎的URL抓取池类
"""

import asyncio
import logging
from typing import Optional
from logger import get_logger

class URLPool:
    def __init__(self) -> None:
        self._queue: asyncio.Queue[str] = asyncio.Queue()
        self._seen_urls: set[str] = set()
        self._lock: asyncio.Lock = asyncio.Lock()
        self.logger = get_logger('url_pool')

    async def add_url(self, url: str) -> None:
        """向池中添加一个URL，如果该URL未被添加过。"""
        async with self._lock:
            if url not in self._seen_urls:
                self._seen_urls.add(url)
                self._queue.put_nowait(url)
                self.logger.debug(f"添加URL到池中: {url}")
            else:
                self.logger.debug(f"URL已存在，跳过: {url}")

    async def get_url(self, timeout: Optional[float] = 5.0) -> Optional[str]:
        """从池中获取一个URL。如果池为空，会等待指定时间后返回None。"""
        try:
            # 使用 asyncio.wait_for 设置超时，避免无限等待
            url = await asyncio.wait_for(self._queue.get(), timeout=timeout)
            self.logger.debug(f"从池中获取URL: {url}")
            return url
        except asyncio.TimeoutError:
            # 超时表示URL池为空，返回None让Worker退出
            # 只在调试模式下记录超时日志，避免正常等待时产生过多日志
            if self.logger.isEnabledFor(logging.DEBUG):
                self.logger.debug("URL池获取超时，池为空")
            return None
        except asyncio.CancelledError:
            # 如果主任务被取消，这里可以返回None或处理取消逻辑
            self.logger.info("URL池获取被取消")
            return None

    def is_empty(self) -> bool:
        """检查URL池是否为空。
        
        注意：asyncio.Queue.empty()的状态是不可靠的，因为队列状态可能在调用
        empty()和后续的get()之间发生变化。建议使用get_url()的超时机制来判断。
        """
        return self._queue.empty()

    def size(self) -> int:
        """返回URL池中URL的数量。"""
        return self._queue.qsize()



